# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import collections
import functools
import logging

from absl import app, flags

from coin.base.timestamp import get_timestamp
from coin.experimental.xguo.fruit.util.util import VolatilityCalculator
from coin.experimental.xguo.fruit.util.record_writer import RecordWriter
from coin.experimental.xguo.fruit.base.single_product import SingleProductPriceCalculator
from coin.experimental.xguo.fruit.base.strategy_base import (
    BaseStrategy,
    BaseStrategyConfig,
)

FLAGS = flags.FLAGS


class AppleStrategyConfig(BaseStrategyConfig):
  pass


class AppleStrategy(BaseStrategy):
  def __init__(self, config):
    super().__init__(config)
    self._vol_calculator = collections.defaultdict(
        functools.partial(VolatilityCalculator, window_sec=30))
    self._price_calculator = {}
    self._vol_writer = RecordWriter('volatility_%s' % get_timestamp())

  def on_book(self, book):
    if book.product not in self._price_calculator:
      self._price_calculator[book.product] = SingleProductPriceCalculator(book.product)
    price_calculator = self._price_calculator[book.product]
    price_calculator.update(book)
    vol_calculator = self._vol_calculator[book.product]
    vol_calculator.update(book.timestamp, price_calculator.midp)
    if vol_calculator.is_ready():
      vol = vol_calculator.get_volatility()
      data = {'timestamp': book.timestamp, 'symbol': book.product.full_symbol, 'volatility': vol}
      self._vol_writer.write(data)

  def on_trade(self, product, trade_event):
    pass


def main(_):
  config = AppleStrategyConfig.from_file(FLAGS.config)
  strategy = AppleStrategy(config)
  strategy.run()


if __name__ == '__main__':
  flags.DEFINE_string('config', None, 'Strategy config file.')
  logging.basicConfig(level='INFO', format='%(levelname)s %(asctime)s %(name)s] %(message)s')
  app.run(main)
